#include "pnp.h"

// solve R, t using Gauss-Newton method
void bundle_adjustment_gauss_newton(const VecVector3d &points_3d,
                                    const VecVector2d &points_2d,
                                    const Mat &K,
                                    Sophus::SE3d &pose)
{
    typedef Eigen::Matrix<double, 6, 1> Vector6d;
    const int iterations = 10;
    double cost = 0, lastCost = 0;
    double fx = K.at<double>(0, 0);
    double fy = K.at<double>(1, 1);
    double cx = K.at<double>(0, 2);
    double cy = K.at<double>(1, 2);

    for (int iter = 0; iter < iterations; iter++)
    {
        Eigen::Matrix<double, 6, 6> H = Eigen::Matrix<double, 6, 6>::Zero();
        Vector6d g = Vector6d::Zero();

        cost = 0;
        // compute cost
        for (int i = 0; i < points_3d.size(); i++)
        {
            Eigen::Vector3d pc = pose * points_3d[i];
            double inv_z = 1.0 / pc[2];
            double inv_z2 = inv_z * inv_z;
            //          X'                 Y'
            // u = fx * -- + cx,  v = fy * -- + cy
            //          Z'                 Z'
            Eigen::Vector2d reproj(fx * pc[0] / pc[2] + cx, fy * pc[1] / pc[2] + cy);

            // reprojection error
            Eigen::Vector2d e = points_2d[i] - reproj;
            cost += e.squaredNorm();

            //       [fx         fx*X'     fx*X'Y'         fx*X'^2     fx*Y']
            //       |--    0  - ----    - -------    fx + -------  - ------|
            //       |Z'         Z'^2       Z'^2             Z'^2      Z'   |
            // J = - |                                                      |
            //       |     fy    fy*Y'       fy*Y'^2     fy*X'Y'      fy*X' |
            //       | 0   --  - ----  -fy - -------     -------      ----  |
            //       [     Z'    Z'^2         Z'^2        Z'^2         Z'   ]
            Eigen::Matrix<double, 2, 6> J;
            J << -fx * inv_z,
                0,
                fx * pc[0] * inv_z2,
                fx * pc[0] * pc[1] * inv_z2,
                -fx - fx * pc[0] * pc[0] * inv_z2,
                fx * pc[1] * inv_z,
                0,
                -fy * inv_z,
                fy * pc[1] * inv_z2,
                fy + fy * pc[1] * pc[1] * inv_z2,
                -fy * pc[0] * pc[1] * inv_z2,
                -fy * pc[0] * inv_z;

            H += J.transpose() * J;
            g += -J.transpose() * e;
        }

        Vector6d dx;
        dx = H.ldlt().solve(g);

        if (isnan(dx[0]))
        {
            cout << "result is nan!" << endl;
            break;
        }

        if (iter > 0 && cost >= lastCost)
        {
            // cost increase, update is not good
            cout << "cost: " << cost << ", last cost: " << lastCost << endl;
            break;
        }

        // update your estimation by left multiplication
        pose = Sophus::SE3d::exp(dx) * pose;
        lastCost = cost;

        cout << "iteration " << iter << " cost=" << std::setprecision(12) << cost << endl;
        if (dx.norm() < 1e-6)
        {
            // converge
            break;
        }
    }

    cout << "pose by g-n: \n"
         << pose.matrix() << endl;
}
